import os
import glob
import yaml
from source.proprecess import *
import streamlit as st
from source.control import param2func


def transform2func(transform_name):
    transform_map = {
        '高斯滤波': gaussian_blur,
        '中值滤波': median_blur,
        '伽马变换': gamma_transform,
        '对比度': adjust_contrast,
        '清晰度': adjust_sharpness,
        '亮度': adjust_brightness,
        '颜色平衡': adjust_color,
        '直方图均衡化': equalize,
        '自适应均衡化': clahe,
        '细节增强': detail_enhance,
        '固定阈值': custom_threshold,
        '自适应阈值': adaptive_threshold,
        '形态学操作': morphology,
        '开运算': morphology_open,
        '闭运算': morphology_close,
        'Sobel': edge_sobel,
        'Laplacian': edge_laplacian,
        'Canny': edge_canney
    }

    return transform_map[transform_name]


def get_images(image_dir):
    assert os.path.isdir(image_dir), "{} is not a directory".format(image_dir)

    images = set()
    image_dir = os.path.abspath(image_dir)
    exts = ['jpg', 'jpeg', 'png', 'bmp']
    exts += [ext.upper() for ext in exts]
    for ext in exts:
        images.update(
            glob.glob('{}/**/*.{}'.format(image_dir, ext), recursive=True))
    images = list(images)

    assert len(images) > 0, "no image found in {}".format(image_dir)
    print("Found {} images in total.".format(len(images)))

    return images


def load_transforms_config(config_file):
    with open(config_file, encoding='utf-8') as f:
        transforms = yaml.safe_load(f)

    return transforms


def select_transforms(transforms, execute_mode):
    # if not execute_mode:
    if execute_mode == '单个变换':
        with st.sidebar.expander("**选择一个变换:**", expanded=True):
            transform_name = st.selectbox(
                "选择一个变换：", list(transforms.keys()), label_visibility='collapsed')
            transform_args = show_transform_control(
                transforms[transform_name], 0)
        return [[transform_name, transform_args]]

    with st.sidebar.expander('**选择变换1**', expanded=True):
        transform_name = st.selectbox(
            '选择变换1', list(transforms.keys()), label_visibility='collapsed')
        transform_args = show_transform_control(transforms[transform_name], 0)

    transform_names = [[transform_name, transform_args]]
    n = 1
    while transform_names[-1][0] != "None":

        with st.sidebar.expander(f'**选择变换{len(transform_names) + 1}**', expanded=True):
            transform_name = st.selectbox(
                f"选择变换{len(transform_names) + 1}",
                ["None"] + list(transforms.keys()),
                label_visibility='collapsed')
            if transform_name != 'None':
                transform_args = show_transform_control(
                    transforms[transform_name], n)
            else:
                transform_args = None

        transform_names.append([transform_name, transform_args])
        n += 1

    return transform_names[:-1]


def show_transform_control(transform_params, n_for_hash):
    param_values = {}
    if len(transform_params) == 0:
        st.sidebar.text("Transform has no parameters")
    else:
        for param in transform_params:
            control_function = param2func[param["type"]]
            if isinstance(param['param_name'], list):
                returned_values = control_function(**param, n_for_hash=n_for_hash)
                for name, value in zip(param['param_name'], returned_values):
                    param_values[name] = value
            else:
                param_values[param['param_name']] = control_function(
                    **param, n_for_hash=n_for_hash)
    return param_values
